import os,sys,pickle,urllib2,time,signal
import subprocess as sub
import threading  as td
from socket import *

from Queue import Queue



def open_port(ip,*port):
    target_ip = ip
    open_ports=[]
    for i in port:
        s = socket(AF_INET, SOCK_STREAM)
        s.settimeout(2)
        result = s.connect_ex((target_ip, i))
        if(result == 0):
            open_ports.append(i)
        try:
            s.close()
        except:
            pass
    return open_ports

class Worker(td.Thread):
    """Thread executing tasks from a given tasks queue"""
    def __init__(self, tasks):
        td.Thread.__init__(self)
        self.tasks = tasks
        self.daemon = True
        self.start()
    
    def run(self):
        while True:
            func, args, kargs = self.tasks.get()
            try: func(*args, **kargs)
            except Exception, e: print e
            self.tasks.task_done()

class ThreadPool:
    """Pool of threads consuming tasks from a queue"""
    def __init__(self, num_threads):
        self.tasks = Queue(num_threads)
        for _ in range(num_threads): Worker(self.tasks)

    def add_task(self, func, *args, **kargs):
        """Add a task to the queue"""
        self.tasks.put((func, args, kargs))

    def wait_completion(self):
        """Wait for completion of all the tasks in the queue"""
        self.tasks.join()
        
class Log(object):
	"""docstring for Log"""
	def __init__(self):
		pass

class ProcessData(object):
	"""docstring for ProcessData"""
	def __init__(self, finput,work_no):
		self.finput=finput
		self.work_no=work_no
		self.input={}
		self.output={}
		tmp=os.popen("mkdir in;mkdir out;mkdir index;mkdir log;mkdir pics")
		os.popen("Xvfb :1 -screen 0 1024x768x24&")
		

	def init(self):
		tmp=os.popen("cat "+self.finput + "| wc -l").read()
		total_count=int(tmp)
		single_count=(total_count/25)+1
		tmp=os.popen("split -l "+str(single_count)+" "+self.finput+" -d -a 2 "+"in/in_")
		pass

	def save(self):
		tmp=os.popen("cat out/out_* > "+self.finput+".res")

	def clean(self):
		tmp=os.popen("rm in/in_*")
		tmp=os.popen("rm out/out_*")


class CmdTimer(td.Thread):
	def __init__(self, cmd, timeout):
        	td.Thread.__init__(self)
        	self.cmd = cmd
        	self.timeout = timeout

	def run(self):
		self.p = sub.Popen(self.cmd)
        	self.p.wait()

    	def exec_cmd(self):
        	self.start()
        	self.join(self.timeout)
        	if self.isAlive():
            		#self.p.terminate()
			os.kill(self.p.pid, signal.SIGKILL)
            		self.join()


class Run(object):
	"""docstring for Run"""
	def __init__(self, finput,no):
		self.no = no
		self.finput=finput
		self.pool = ThreadPool(self.no)
		self.procdata=ProcessData(finput,self.no)

	def start(self):
		self.procdata.init()
		d=Crawl()
		for i in range(0,25):
			t=["0"+str(i) if i<10 else str(i)][0]
			finput="in/in_"+t
			foutput="out/out_"+t
			d.pre_process(finput, self.procdata.input)
			for v in self.procdata.input.keys():
				
				self.pool.add_task(d.task,v,self.procdata.output)
			self.pool.wait_completion()
			f=open(foutput,"w")
			f.writelines(self.procdata.output.keys())
			self.procdata.output.clear()
			self.procdata.input.clear()

			f.close()
			print finput,"finish"

		self.procdata.save()
		self.procdata.clean()

"""
port 1=80:0 or 1
port 2=443:0 or 1
dns_reverse:  1,domain or 0,0
80_redirect: 1,url or 0,0
80_reverse: 1,url or 0,0
443_redirect: 1,url or 0,0
443_reverse: 1,url or 0,0


"""

class Crawl(object):
	"""docstring for Crawl"""
	def __init__(self):
		super(Crawl, self).__init__()
		pass

	def pre_process(self,finput,obuffer):
		for line in open(finput,"r"):
			line=line.strip("\n")
			#print line
			obuffer[line]=1
		pass


	def task(self,ip,obuffer):
		r=self.scan_ip(ip)
		obuffer[ip+"#"+str(r)+"\n"]=1
		pass

	def download_page(self,ip):
		try_time=1
		log_name="log/"+str(ip)+".log"
		page_name="index/"+str(ip)+".html"
		cmd=["wget","--connect-timeout","15","-t",str(try_time),str(ip),"-o",log_name,"-O",page_name]
		CmdTimer(cmd, 30).exec_cmd()
		#c.run()

	def get_screenshot(self,ip):
		os.putenv("DISPLAY","localhost:1.0")
		url="http://"+ip
		#os.popen("SCREENSHOOT/wkhtmltoimage-i386 --use-xserver --height 768 --width 1024 "+url+" PICS/"+ip+".png")
		cmd=["./wkhtmltoimage-i386","--use-xserver","--height","768","--width","1024",url,"pics/"+ip+".png"]
		CmdTimer(cmd, 90).exec_cmd()


	def scan_ip(self,ip):
		print ip
		ports=open_port(ip,80)

		if not ports:
			return 0
		else:
			if 80 in ports:
				try:
					self.download_page(ip)
					self.get_screenshot(ip)
					return 1
				except:
					f=open("except.log","a")
					f.write(ip+"\n")
					f.close()
					return 0
					pass



if __name__ == '__main__':
	finput=sys.argv[1]
	workno=sys.argv[2]
	r=Run(finput,int(workno))
	r.start()
	#c=Crawl()
	##p=ProcessData("in_00",25)
	#p.init()
	#ip="54.243.104.3"	
	#print c.reverse_domain("204.236.195.189")
	#ip="204.236.195.85"
	#print c.reverse_domain(ip)

	
